#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import subprocess
import sys
import time
import logging
from typing import List, Tuple, Optional, Dict, Any
from dataclasses import dataclass
from pathlib import Path


@dataclass
class ClaudeConfig:
    """Configuration class for Claude CLI settings."""
    timeout: int = 300
    max_history_entries: int = 50
    claude_command: str = "claude"
    claude_args: List[str] = None
    log_level: str = "INFO"
    retry_attempts: int = 3
    retry_delay: float = 1.0

    def __post_init__(self):
        if self.claude_args is None:
            self.claude_args = ["--print", "--dangerously-skip-permissions"]


class ClaudeError(Exception):
    """Base exception for Claude CLI operations."""
    pass


class ClaudeTimeoutError(ClaudeError):
    """Raised when Claude CLI operation times out."""
    pass


class ClaudeProcessError(ClaudeError):
    """Raised when Claude CLI process fails."""
    def __init__(self, returncode: int, cmd: List[str], stderr: str = None):
        self.returncode = returncode
        self.cmd = cmd
        self.stderr = stderr
        super().__init__(f"Command {' '.join(cmd)} failed with return code {returncode}")


class OutputFormatter:
    """Handles formatting and display of Claude CLI output."""

    @staticmethod
    def format_task_start(prompt: str, timeout: int) -> str:
        return f"⚙️ Calling claude CLI with task instruction...\n📝 Task: {prompt}\n⏱️ Timeout: {timeout}s\n{'-' * 60}"

    @staticmethod
    def format_success() -> str:
        return f"✅ Claude CLI executed successfully.\n{'-' * 60}"

    @staticmethod
    def format_error(error_type: str, message: str) -> str:
        return f"❌ Error: {error_type}\n{message}"


class ClaudeCLI:
    """
    An interactive wrapper class for calling Claude CLI with full tool support.
    """

    def __init__(self, config: Optional[ClaudeConfig] = None, logger: Optional[logging.Logger] = None):
        """
        Initialize ClaudeCLI with configuration options.

        Args:
            config: Configuration object with timeout, retry settings, etc.
            logger: Optional logger instance for custom logging
        """
        self.config = config or ClaudeConfig()
        self.logger = logger or self._setup_logger()
        self.process = None
        self._timeout_occurred = False
        self._session_active = False
        self._conversation_history: List[Tuple[str, str]] = []
        self.formatter = OutputFormatter()

    def _setup_logger(self) -> logging.Logger:
        """Setup default logger for Claude CLI operations."""
        logger = logging.getLogger("claude_cli")
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
            logger.setLevel(getattr(logging, self.config.log_level))
        return logger

    def run(self, prompt: str, verbose: bool = True) -> str:
        """
        Execute a task using Claude CLI with retry logic and improved error handling.

        Args:
            prompt: The task instruction to send to Claude CLI
            verbose: Whether to print progress information

        Returns:
            str: Claude's response output

        Raises:
            ClaudeTimeoutError: If execution times out
            ClaudeProcessError: If Claude CLI process fails
            ClaudeError: For other Claude-related errors
        """
        if verbose:
            print(self.formatter.format_task_start(prompt, self.config.timeout))

        self.logger.info(f"Starting Claude CLI execution with timeout {self.config.timeout}s")

        last_exception = None
        for attempt in range(self.config.retry_attempts):
            try:
                if attempt > 0:
                    self.logger.info(f"Retry attempt {attempt + 1}/{self.config.retry_attempts}")
                    if verbose:
                        print(f"🔄 Retry attempt {attempt + 1}/{self.config.retry_attempts}")
                    time.sleep(self.config.retry_delay * attempt)

                result = self._execute_claude_command(prompt)

                if verbose:
                    print(result.stdout, flush=True)
                    print(self.formatter.format_success())

                self.logger.info("Claude CLI execution completed successfully")
                return result.stdout

            except subprocess.TimeoutExpired as e:
                last_exception = ClaudeTimeoutError(f"Claude CLI execution timed out after {self.config.timeout}s")
                self.logger.warning(f"Attempt {attempt + 1} timed out")
                if verbose:
                    print(f"⏰ Timeout on attempt {attempt + 1}")

            except subprocess.CalledProcessError as e:
                last_exception = ClaudeProcessError(e.returncode, e.cmd, e.stderr)
                self.logger.error(f"Attempt {attempt + 1} failed with return code {e.returncode}: {e.stderr}")
                if verbose:
                    print(f"🔥 Error on attempt {attempt + 1}: {e.stderr}")

            except FileNotFoundError as e:
                error_msg = f"'{self.config.claude_command}' command not found. Please ensure Claude CLI is installed."
                self.logger.error(error_msg)
                if verbose:
                    print(self.formatter.format_error("Command Not Found", error_msg))
                raise ClaudeError(error_msg) from e

            except KeyboardInterrupt as e:
                self.logger.info("User interrupted execution")
                if verbose:
                    print("⚠️  User interrupted the execution")
                raise

        # All retry attempts failed
        error_msg = f"All {self.config.retry_attempts} attempts failed"
        self.logger.error(error_msg)
        if verbose:
            print(self.formatter.format_error("Max Retries Exceeded", error_msg))
        raise last_exception or ClaudeError(error_msg)

    def _execute_claude_command(self, prompt: str) -> subprocess.CompletedProcess:
        """Execute the Claude CLI command with given prompt."""
        cmd = [self.config.claude_command] + self.config.claude_args + [prompt]

        self.logger.debug(f"Executing command: {' '.join(cmd)}")

        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=self.config.timeout
        )

        if result.returncode != 0:
            raise subprocess.CalledProcessError(
                result.returncode,
                cmd,
                stderr=result.stderr
            )

        return result

    def start_session(self, verbose: bool = True) -> None:
        """
        Start a conversation session with optimized history management.

        Args:
            verbose: Whether to print progress information
        """
        if self._session_active:
            self.logger.warning("Session already active")
            if verbose:
                print("⚠️  Session already active")
            return

        if verbose:
            print("🚀 Starting Claude conversation session...")
            print("-" * 60)

        self._session_active = True
        self._conversation_history = []
        self.logger.info("Claude conversation session started")

        if verbose:
            print("✅ Session started successfully")

    def send_message(self, message: str, verbose: bool = True) -> str:
        """
        Send a message in the conversation session with optimized history management.

        Args:
            message: The message to send to Claude
            verbose: Whether to print progress information

        Returns:
            str: Claude's response

        Raises:
            RuntimeError: If no active session
        """
        if not self._session_active:
            raise RuntimeError("No active session. Call start_session() first.")

        if verbose:
            print(f"📝 Sending: {message}")
            print("-" * 60)

        self.logger.info(f"Sending message in session (history length: {len(self._conversation_history)})")

        # Build efficient context from conversation history
        context_prompt = self._build_context_prompt(message)

        # Use the working run() method with context
        response = self.run(context_prompt, verbose=verbose)

        # Store in conversation history with automatic cleanup
        self._add_to_history(message, response)

        return response

    def _build_context_prompt(self, current_message: str) -> str:
        """
        Build context prompt efficiently, managing history size.

        Args:
            current_message: The current user message

        Returns:
            str: Formatted context prompt
        """
        if not self._conversation_history:
            return current_message

        # Use only recent history to prevent prompt bloat
        recent_history = self._get_recent_history()

        context_parts = ["Previous conversation:"]
        for user_msg, claude_response in recent_history:
            context_parts.append(f"User: {user_msg}")
            context_parts.append(f"Claude: {claude_response}")
            context_parts.append("")  # Empty line for separation

        context_parts.append(f"User: {current_message}")
        return "\n".join(context_parts)

    def _get_recent_history(self) -> List[Tuple[str, str]]:
        """
        Get recent conversation history within configured limits.

        Returns:
            List of recent (user_message, claude_response) tuples
        """
        max_entries = self.config.max_history_entries
        if len(self._conversation_history) <= max_entries:
            return self._conversation_history

        # Keep most recent entries and log when truncating
        self.logger.info(f"Truncating history from {len(self._conversation_history)} to {max_entries} entries")
        return self._conversation_history[-max_entries:]

    def _add_to_history(self, message: str, response: str) -> None:
        """
        Add message and response to history with automatic cleanup.

        Args:
            message: User message
            response: Claude's response
        """
        self._conversation_history.append((message, response))

        # Auto-cleanup if history gets too long
        if len(self._conversation_history) > self.config.max_history_entries * 2:
            self._cleanup_history()

    def _cleanup_history(self) -> None:
        """
        Clean up conversation history to maintain performance.
        """
        target_size = self.config.max_history_entries
        if len(self._conversation_history) > target_size:
            removed_count = len(self._conversation_history) - target_size
            self._conversation_history = self._conversation_history[-target_size:]
            self.logger.info(f"Cleaned up history: removed {removed_count} old entries")

    def end_session(self, verbose: bool = True) -> None:
        """
        End the active Claude session and clean up resources.

        Args:
            verbose: Whether to print progress information
        """
        if self._session_active:
            if verbose:
                print("🔚 Ending Claude session...")

            history_length = len(self._conversation_history)
            self.logger.info(f"Ending session with {history_length} conversation entries")

            self._session_active = False
            self._conversation_history.clear()  # More explicit cleanup

            if verbose:
                print("✅ Session ended successfully")
        else:
            self.logger.warning("Attempted to end session, but no active session found")

    def get_session_stats(self) -> Dict[str, Any]:
        """
        Get statistics about the current session.

        Returns:
            Dict containing session statistics
        """
        return {
            "session_active": self._session_active,
            "history_length": len(self._conversation_history),
            "max_history_entries": self.config.max_history_entries,
            "timeout": self.config.timeout,
            "retry_attempts": self.config.retry_attempts
        }

    def _handle_timeout(self) -> None:
        """Handle timeout by setting flag and terminating process."""
        self._timeout_occurred = True
        self.logger.warning(f"Timeout reached ({self.config.timeout}s), terminating process")
        print(f"⏰ Timeout reached ({self.config.timeout}s), terminating process...")
        if self.process:
            self._cleanup_process()

    def _cleanup_process(self) -> None:
        """Clean up process resources with proper error handling."""
        if self.process:
            try:
                # Close stdin first to signal end of input
                if self.process.stdin:
                    self.process.stdin.close()

                if self.process.poll() is None:
                    self.process.terminate()
                    self.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.logger.warning("Process termination timed out after 5 seconds")
                if self.process.poll() is None:
                    try:
                        self.process.kill()
                        self.process.wait(timeout=2)
                    except (subprocess.TimeoutExpired, ProcessLookupError):
                        self.logger.error("Failed to kill process forcefully")
            except (ProcessLookupError, OSError) as e:
                self.logger.info(f"Process cleanup completed: {e}")
            except Exception as e:
                self.logger.error(f"Unexpected error during process cleanup: {e}")
            finally:
                self.process = None
                self._session_active = False


def demo_single_execution() -> None:
    """
    Demo function showing single-shot Claude CLI execution.
    """
    print("\n🚀 Testing single execution mode...")
    config = ClaudeConfig(timeout=60, log_level="INFO")
    cli = ClaudeCLI(config=config)

    try:
        response = cli.run("What is 2+2? Respond briefly.")
        print(f"\n📝 Single execution result: {len(response)} characters")
    except ClaudeError as e:
        print(f"❌ Claude error: {e}")
    except Exception as e:
        print(f"❌ Unexpected error: {e}")


def demo_conversation_session() -> None:
    """
    Demo function showing conversation session functionality.
    """
    print("\n🚀 Testing conversation session...")
    config = ClaudeConfig(timeout=60, max_history_entries=10, log_level="INFO")
    cli = ClaudeCLI(config=config)

    try:
        cli.start_session()

        # Test conversation continuity
        print("\n📤 Testing conversation flow...")
        response1 = cli.send_message("Say hello and what is 2+2?")
        response2 = cli.send_message("What was my previous question?")
        response3 = cli.send_message("Create a Python function that adds two numbers")

        # Show session stats
        stats = cli.get_session_stats()
        print(f"\n📊 Session stats: {stats}")

        cli.end_session()

        print("\n✅ Conversation test completed successfully!")
        print(f"📝 Response lengths: {len(response1)}, {len(response2)}, {len(response3)} characters")

    except ClaudeError as e:
        print(f"❌ Claude error: {e}")
        cli.end_session(verbose=False)
    except KeyboardInterrupt:
        print("\n⚠️  User interrupted")
        cli.end_session(verbose=False)
    except Exception as e:
        print(f"❌ Unexpected error: {e}")
        cli.end_session(verbose=False)


def main() -> None:
    """
    Main demo function showcasing ClaudeCLI capabilities.
    """
    print("🚀 Starting Claude CLI Demo")
    print("=" * 60)

    # Configure logging for demo
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )

    try:
        # Demo single execution
        demo_single_execution()

        # Demo conversation session
        demo_conversation_session()

        print("\n✅ All demos completed successfully!")

    except KeyboardInterrupt:
        print("\n⚠️  Demo interrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"❌ Demo failed with unexpected error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()